import numpy as np
import os
import re
from os import listdir
from os.path import join

import torch
from scipy import io
import pandas as pd
# from torch.utils.data import DataLoader, Dataset
from datasets.thu_chl.process import *

repr_map = {'eventFrame': get_eventFrame,
            'eventAccuFrame': get_eventAccuFrame,
            'timeSurface': get_timeSurface,
            'eventCount': get_eventCount}


def events_to_voxel_grid_pytorch(events, num_bins, width, height, device):
    """
    Build a voxel grid with bilinear interpolation in the time domain from a set of events.
    :param events: a [N x 4] NumPy array containing one event per row in the form: [timestamp, x, y, polarity]
    :param num_bins: number of bins in the temporal axis of the voxel grid
    :param width, height: dimensions of the voxel grid
    :param device: device to use to perform computations
    :return voxel_grid: PyTorch event tensor (on the device specified)
    """

    assert (events.shape[1] == 4)
    assert (num_bins > 0)
    assert (width > 0)
    assert (height > 0)

    with torch.no_grad():

        events_torch = torch.from_numpy(events).float()
        events_torch = events_torch.to(device)

        voxel_grid = torch.zeros(num_bins, height, width, dtype=torch.float32, device=device)
        if events_torch.shape[0] == 0:
            return voxel_grid

        voxel_grid = voxel_grid.flatten()

        # normalize the event timestamps so that they lie between 0 and num_bins
        last_stamp = events_torch[-1, 0]
        first_stamp = events_torch[0, 0]
        deltaT = float(last_stamp - first_stamp)

        if deltaT == 0:
            deltaT = 1.0

        events_torch[:, 0] = (num_bins - 1) * (events_torch[:, 0] - first_stamp) / deltaT
        ts = events_torch[:, 0]
        xs = events_torch[:, 1].long()
        ys = events_torch[:, 2].long()
        pols = events_torch[:, 3].float()
        pols[pols == 0] = -1  # polarity should be +1 / -1

        tis = torch.floor(ts)
        tis_long = tis.long()
        dts = ts - tis
        vals_left = pols * (1.0 - dts.float())
        vals_right = pols * dts.float()

        valid_indices = tis < num_bins
        valid_indices &= tis >= 0
        voxel_grid.index_add_(dim=0,
                              index=xs[valid_indices] + ys[valid_indices]
                                    * width + tis_long[valid_indices] * width * height,
                              source=vals_left[valid_indices])

        valid_indices = (tis + 1) < num_bins
        valid_indices &= tis >= 0

        voxel_grid.index_add_(dim=0,
                              index=xs[valid_indices] + ys[valid_indices] * width
                                    + (tis_long[valid_indices] + 1) * width * height,
                              source=vals_right[valid_indices])

        voxel_grid = voxel_grid.view(num_bins, height, width)

    return voxel_grid


def events_to_voxel_all(events, seq_len, num_bins, width, height, device):
    voxel_len = seq_len * num_bins
    voxel_grid_all = np.zeros((num_bins * seq_len, 1, height, width))
    voxel_grid = events_to_voxel_grid_pytorch(events, voxel_len, width, height, device)
    voxel_grid = voxel_grid.unsqueeze(1).cpu().numpy()
    voxel_grid_all[:voxel_len] = voxel_grid
    return voxel_grid_all


# left or right move all event locations randomly
def random_shift_events(events, max_shift=20, resolution=(180, 240)):
    H, W = resolution
    x_shift, y_shift = np.random.randint(-max_shift, max_shift + 1, size=(2,))
    events[:, 0] += x_shift
    events[:, 1] += y_shift

    valid_events = (events[:, 0] >= 0) & (events[:, 0] < W) & (events[:, 1] >= 0) & (events[:, 1] < H)
    events = events[valid_events]

    return events


# flip half of the event images along the x dimension
def random_flip_events_along_x(events, resolution=(180, 240), p=0.5):
    H, W = resolution
    if np.random.random() < p:
        events[:, 0] = W - 1 - events[:, 0]
    return events


class THU_EACT_50_CHL:
    def __init__(self, datafile="../THU-EACT-50-CHL", train=True, augmentation=False,
                 repr=['timeSurface', 'eventFrame'], time_num=16, ret_file_name=False,
                 frame_transform=None, label_transform=None, voxel=True):
        list_file_name = join(datafile, "train.txt") if train else join(datafile, "test.txt")

        self.files = []
        self.labels = []
        self.augmentation = augmentation
        self.datafile = datafile
        self.phase = train
        self.repr = repr
        self.time_num = time_num
        self.length = 4
        self.ret_file_name = ret_file_name
        self.voxel = voxel
        self.frame_transform = frame_transform
        self.label_transform = label_transform

        list_file = open(list_file_name, "r")
        for line in list_file:
            file, label = line.split(" ")
            self.files.append(file)
            self.labels.append(int(label))
        list_file.close()

        self.classes = np.unique(self.labels)

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        """
        returns events and label, loading events from aedat
        :param idx:
        :return: x,y,t,p,  label
        events: (x,y,t,p)  x:[0,346] y:[0,260] p:0/1
        """
        label = self.labels[idx]
        f = self.files[idx]
        f = f.split('DVS-action-data-npy/')[-1]
        f = os.path.join(self.datafile, f)

        events_input = np.load(f).astype(np.float32)
        if self.voxel:
            # 指定的顺序
            desired_order = [2, 0, 1, 3]  # 这表示希望将第二维按照指定的顺序重新排列
            # 使用指定的顺序重新排列数组
            events_input = events_input[:, desired_order]  # (x,y,t,p)->(t,x,y,p)
            event_voxel_low = events_to_voxel_all(events_input, self.length, self.time_num // self.length,
                                                  346, 260,
                                                  device='cpu')  # # (30*num_bins[1],1, 96, 96)
            event_voxel_high = events_to_voxel_all(events_input, self.length, self.time_num // self.length,
                                                   346, 260,
                                                   device='cpu')  # # (30*num_bins[1],1, 96, 96)
            if self.phase:
                event_voxel_low, event_voxel_high = RandomCrop(event_voxel_low, event_voxel_high, (256, 256))
                event_voxel_low, event_voxel_high = HorizontalFlip(event_voxel_low, event_voxel_high)
            else:
                event_voxel_low, event_voxel_high = CenterCrop(event_voxel_low, event_voxel_high, (256, 256))

            reprs = np.concatenate((event_voxel_low, event_voxel_high), axis=1)  # (T,C, 96, 96)
        else:
            # normalize the timestamps
            _min = events_input[:, 2].min()
            _max = events_input[:, 2].max()
            events_input[:, 2] = (events_input[:, 2] - _min) / (_max - _min)

            if self.augmentation:
                events_input = random_shift_events(events_input)
                events_input = random_flip_events_along_x(events_input)

            reprs = []
            for repr_name in self.repr:
                # 输入(t,x,y,p), 输出(T,y,x) y:[0,260] x:[0:346]
                repr_array = repr_map[repr_name](events_input[:, 2], events_input[:, 0].astype(np.int32),
                                                 events_input[:, 1].astype(np.int32),
                                                 events_input[:, 3],
                                                 repr_size=(260, 346), time_num=self.time_num)

                # standardization
                mu = np.mean(repr_array)
                sigma = np.std(repr_array)
                repr_array = (repr_array - mu) / sigma

                reprs.append(repr_array)

            reprs = np.array(reprs)  # (C,T,260,346)
            reprs = np.transpose(reprs, (1, 0, 3, 2))  # (T,C,346,260)

        if self.frame_transform is not None:
            reprs = self.frame_transform(reprs)

        if self.label_transform is not None:
            label = self.label_transform(label)

        if self.ret_file_name:
            # file_name = re.findall(r'A[\w-]+', f)[0]
            file_name = f.split('/')[-1].split('.')[0]
            return reprs, label, file_name
        else:
            return reprs, label
